{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "W4_Homework.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/CIS-522/course-content/blob/main/tutorials/W04_Optimization/W4_Homework.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KNU5GiXY_Dq6"
      },
      "source": [
        "# CIS-522 Week 4 Homework\n",
        "\n",
        "\n",
        "**Instructor:** Lyle Ungar\n",
        "\n",
        "**Content Creator:** Rongguang Wang, Jordan Lei"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "szKIY9I-_Qa-"
      },
      "source": [
        "---\n",
        "## Preface\n",
        "Since this week's homework requires coding, we recommend saving this notebook in your google Drive (`File -> Save a copy in Drive`), and share the link to the final version in the subscription airtable form. You can also attach the code to the form if you prefer off-colab coding."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vAdEW6Eq-8IN",
        "cellView": "form"
      },
      "source": [
        "#@markdown What is your Pennkey and pod? (text, not numbers, e.g. bfranklin)\n",
        "my_pennkey = '' #@param {type:\"string\"}\n",
        "my_pod = 'Select' #@param ['Select', 'euclidean-wombat', 'sublime-newt', 'buoyant-unicorn', 'lackadaisical-manatee','indelible-stingray','superfluous-lyrebird','discreet-reindeer','quizzical-goldfish','ubiquitous-cheetah','nonchalant-crocodile','fashionable-lemur','spiffy-eagle','electric-emu','quotidian-lion','astute-jellyfish', 'quantum-herring']\n",
        "\n",
        "# start timing\n",
        "import time\n",
        "try:t0;\n",
        "except NameError: t0 = time.time()\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3tBIYJK0vr6z"
      },
      "source": [
        "# Setup"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Rq1asJ9vvseM"
      },
      "source": [
        "# imports\n",
        "from __future__ import print_function\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.optim as optim\n",
        "from torchvision import datasets, transforms\n",
        "from torch.optim.lr_scheduler import StepLR\n",
        "from torch.utils.data import Dataset\n",
        "import time\n",
        "import math\n",
        "import numpy as np\n",
        "from bisect import bisect_right\n",
        "import matplotlib.pyplot as plt\n",
        "from collections import defaultdict"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5EZxszBFvvf8",
        "cellView": "form"
      },
      "source": [
        "# @title Figure Settings\n",
        "import ipywidgets as widgets\n",
        "%matplotlib inline \n",
        "fig_w, fig_h = (8, 6)\n",
        "plt.rcParams.update({'figure.figsize': (fig_w, fig_h)})\n",
        "%config InlineBackend.figure_format = 'retina'\n",
        "SMALL_SIZE = 12\n",
        "\n",
        "plt.rcParams.update(plt.rcParamsDefault)\n",
        "plt.rc('font', size=SMALL_SIZE)          # controls default text sizes\n",
        "plt.rc('axes', titlesize=SMALL_SIZE)     # fontsize of the axes title\n",
        "plt.rc('axes', labelsize=SMALL_SIZE)    # fontsize of the x and y labels\n",
        "plt.rc('xtick', labelsize=SMALL_SIZE)    # fontsize of the tick labels\n",
        "plt.rc('ytick', labelsize=SMALL_SIZE)    # fontsize of the tick labels\n",
        "plt.rc('legend', fontsize=SMALL_SIZE)    # legend fontsize\n",
        "plt.rc('figure', titlesize=SMALL_SIZE)  # fontsize of the figure title"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AfbGoVXFA8xG"
      },
      "source": [
        "# Part 1: Paper Reading\n",
        "\n",
        "Please read the paper [ADAM: A METHOD FOR STOCHASTIC OPTIMIZATION](https://arxiv.org/pdf/1412.6980.pdf) and answer the following questions. \n",
        "Note that the paper is in some ways beyond the scope of this course. You can skip most of the details (and all of the proofs) in the paper. Just look at it enough to answer the following questions."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dRE5_s_NxBg_",
        "cellView": "form"
      },
      "source": [
        "#@markdown **Student response**: The paper says that “Some of Adam’s advantages are that, 1) its stepsizes are approximately bounded by the stepsize hyperparameter, 2) it does not require a stationary objective, 3)it works with sparse gradients, and 4) it naturally performs a form of step size annealing“. Please rephrase each of these claims,  to make sure you understand what they mean.\n",
        "rephrase_1 = '' #@param {type:\"string\"}\n",
        "rephrase_2 = '' #@param {type:\"string\"}\n",
        "rephrase_3 = '' #@param {type:\"string\"}\n",
        "rephrase_4 = '' #@param {type:\"string\"}\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8XodjWcyAHGO",
        "cellView": "form"
      },
      "source": [
        "#@markdown **Student response**: In the typical case, what is the bound on the update to the weights, \\theta? (Hint: it’s just a simple hyperparameter)\n",
        "bound = '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "K8R1a2SsyIkd",
        "cellView": "form"
      },
      "source": [
        "#@markdown **Student response**: What is “initialization bias correction”, and why is it needed?\n",
        "bias_correction = '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "F37DryH3yRUI",
        "cellView": "form"
      },
      "source": [
        "#@markdown **Student response**: In English,  what is the “regret” in equation 5, and how does the regret scale with the number of stochastic gradient descent steps?\n",
        "regret = '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "InXIZ48PAyXr"
      },
      "source": [
        "# Part 2: Cyclic Cosine Annealing\n",
        "\n",
        "We have introduced many different variants of gradient descent for global minimum search such as Momentum, Adagrad, RMSProp, and Adam. These optimizers all follow their own pattern of finding this minimum loss by giving rules on how to adjust their step size, but what if we get stuck in a local minimum of a loss landscape? Even though some optimizers try and use ideas of momentum to have the ability to swing out of a local minimum, they are not always that successful.\n",
        "\n",
        "<p align=\"center\">\n",
        "  <img width=\"500\" height=\"200\" src=\"https://miro.medium.com/max/1400/1*9Fca3kpx3pVW8SaYz2pjpw.png\">\n",
        "</p>\n",
        "\n",
        "For example, the model in the left figure slowly descends into a local minimum, while the model in the right figure jumps in and out of several local minimums, searching for a more stable one.\n",
        "\n",
        "One possible solution is cyclic cosine annealing. Instead of trying to add various forms of learning rate decay, let’s reset our learning rate every fixed iterations so that we might be able to more easily pop out of a local minimum. Specifically, we simulate a restart of SGD once $T_i$ epochs are performed, where $i$ is the index of the run. The learning rate update rule within the $i$-th run is:\n",
        "$$\\eta_t=\\eta^i_{min}+\\frac{1}{2}(\\eta^i_{max}-\\eta^i_{min})(1+\\cos(\\frac{T_{cur}}{T_i}\\pi))$$\n",
        "where\n",
        "*   $\\eta^i_{min}$ and $\\eta^i_{max}$ are the range of the learning rate;\n",
        "*   $T_{cur}$ is the count of epoch since last restart;\n",
        "*   $T_i$ is the width between the two restarts.\n",
        "\n",
        "<p align=\"center\">\n",
        "  <img width=\"300\" height=\"200\" src=\"https://miro.medium.com/max/1400/1*GjdEoZcEcQF19JS1y8nVYw.png\">\n",
        "</p>\n",
        "\n",
        "The above figure shows an example of cyclic consine annealing. For more details, check out this [paper](https://arxiv.org/pdf/1608.03983.pdf).\n",
        "\n",
        "You will implement a cyclic cosine annealing learning rate schedueler to perform digit classification. Towards this goal, you have to perform the following steps:\n",
        "\n",
        "1. Define a cyclic consine annealing function which restarts the learning rate at two milestones (5 and 10 in this case). Insert your code in the cell below.\n",
        "\n",
        "   Hint: Modify the learning rate using below code.\n",
        "   ```python\n",
        "    for param_group in optimizer.param_groups: \n",
        "        param_group[’lr’] = lr\n",
        "   ```"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2MK78rlqEO9_"
      },
      "source": [
        "def cyclic_cosine_annealing(optimizer, milestones=[5,10], eta_min=1e-6, last_epoch=0):\n",
        "    ####################################################################\n",
        "    # Fill in missing code below (...),\n",
        "    # then remove or comment the line below to test your function\n",
        "    raise NotImplementedError(\"Add the missing code\")\n",
        "    ####################################################################  \n",
        "    \n",
        "\n",
        "    \n",
        "    ####################################################################  "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XLWGeTV-eM7W"
      },
      "source": [
        "2. Check the learning rate computed by the scheduler you just implemented using the code below. Make sure that you observe behavior similar to that in the above figure. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KvdCHlMHdJCj"
      },
      "source": [
        "model = torch.nn.Linear(2, 1)\n",
        "optimizer = optim.SGD(model.parameters(), lr=1e-2)\n",
        "\n",
        "lr_list = []\n",
        "for i in range(20):\n",
        "    cyclic_cosine_annealing(optimizer, last_epoch=i)\n",
        "    lr_list.append(optimizer.param_groups[0]['lr'])\n",
        "    optimizer.step()\n",
        "\n",
        "with plt.xkcd():\n",
        "    plt.plot(range(20), lr_list, label='Cyclic Consine')\n",
        "    plt.ylabel('Learning Rate')\n",
        "    plt.xlabel('Epoch')\n",
        "    plt.legend()\n",
        "    plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pXhPaacie6kR"
      },
      "source": [
        "3. Train a digit classifier on MNIST using your cyclic cosine annealing scheduler. We provide the training code below."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Il51GVhZfU7E",
        "cellView": "form"
      },
      "source": [
        "# @title MNIST Classifier (run me)\n",
        "class Net(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(Net, self).__init__()\n",
        "        self.dropout = nn.Dropout(0.5)\n",
        "        self.fc1 = nn.Linear(784, 128)\n",
        "        self.fc2 = nn.Linear(128, 10)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = torch.flatten(x, 1)\n",
        "        x = self.fc1(x)\n",
        "        x = F.relu(x)\n",
        "        x = self.dropout(x)\n",
        "        x = self.fc2(x)\n",
        "        output = F.log_softmax(x, dim=1)\n",
        "        return output\n",
        "\n",
        "def train(args, model, device, train_loader, optimizer, epoch):\n",
        "    model.train()\n",
        "    avg_loss = 0.\n",
        "    for batch_idx, (data, target) in enumerate(train_loader):\n",
        "        data, target = data.to(device), target.to(device)\n",
        "        optimizer.zero_grad()\n",
        "        output = model(data)\n",
        "        loss = F.nll_loss(output, target)\n",
        "        avg_loss += loss.item()\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "        if batch_idx % args['log_interval'] == 0:\n",
        "            print('Train Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}'.format(\n",
        "                epoch, batch_idx * len(data), len(train_loader.dataset),\n",
        "                100. * batch_idx / len(train_loader), loss.item()))\n",
        "    avg_loss /= len(train_loader.dataset)\n",
        "    return avg_loss\n",
        "\n",
        "def test(model, device, test_loader):\n",
        "    model.eval()\n",
        "    test_loss = 0\n",
        "    correct = 0\n",
        "    with torch.no_grad():\n",
        "        for data, target in test_loader:\n",
        "            data, target = data.to(device), target.to(device)\n",
        "            output = model(data)\n",
        "            test_loss += F.nll_loss(output, target, reduction='sum').item()  # sum up batch loss\n",
        "            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability\n",
        "            correct += pred.eq(target.view_as(pred)).sum().item()\n",
        "\n",
        "    test_loss /= len(test_loader.dataset)\n",
        "\n",
        "    print('\\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.4f}%)\\n'.format(\n",
        "        test_loss, correct, len(test_loader.dataset),\n",
        "        100. * correct / len(test_loader.dataset)))\n",
        "    return 1. - correct / len(test_loader.dataset)\n",
        "\n",
        "def main(args):\n",
        "    # gpu configuration\n",
        "    use_cuda = not args['no_cuda'] and torch.cuda.is_available()\n",
        "    torch.manual_seed(args['seed'])\n",
        "    device = torch.device('cuda' if use_cuda else 'cpu')\n",
        "    # data loader configuration\n",
        "    train_kwargs = {'batch_size': args['batch_size']}\n",
        "    test_kwargs = {'batch_size': args['test_batch_size']}\n",
        "    if use_cuda:\n",
        "        cuda_kwargs = {'num_workers': 1,\n",
        "                       'pin_memory': True,\n",
        "                       'shuffle': True}\n",
        "        train_kwargs.update(cuda_kwargs)\n",
        "        test_kwargs.update(cuda_kwargs)\n",
        "    # load MNIST dataset\n",
        "    transform=transforms.Compose([\n",
        "        transforms.ToTensor(),\n",
        "        transforms.Normalize((0.1307,), (0.3081,))\n",
        "        ])\n",
        "    train_loader = torch.utils.data.DataLoader(datasets.MNIST('../data', train=True, download=True,\n",
        "                       transform=transform),**train_kwargs)\n",
        "    test_loader = torch.utils.data.DataLoader(datasets.MNIST('../data', train=False,\n",
        "                       transform=transform), **test_kwargs)\n",
        "    # init models\n",
        "    model = Net().to(device)\n",
        "    # config optimizer\n",
        "    optimizer = optim.SGD(model.parameters(), lr=args['lr'], momentum=args['momentum'])\n",
        "    # training and testing\n",
        "    train_list, test_list = [], []\n",
        "    for epoch in range(1, args['epochs'] + 1):\n",
        "        if args['scheduler']:\n",
        "            cyclic_cosine_annealing(optimizer, last_epoch=epoch-1)\n",
        "        loss = train(args, model, device, train_loader, optimizer, epoch) # train\n",
        "        train_list.append(loss)\n",
        "        error = test(model, device, test_loader) # test\n",
        "        test_list.append(error)\n",
        "\n",
        "    return train_list, test_list"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7G-8Abhvh0hX"
      },
      "source": [
        "4. Run the below code to train the network.  Compare the training and testing error of SGD and SGD with cyclic cosine annealing learning rate."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_s9mnfxOiTl2",
        "cellView": "form"
      },
      "source": [
        "# @title Train and Plot (run me)\n",
        "# Training settings\n",
        "args = {'batch_size': 64,\n",
        "        'test_batch_size': 1000,\n",
        "        'epochs': 20,\n",
        "        'lr': 0.01,\n",
        "        'scheduler': False,\n",
        "        'momentum': 0.9,\n",
        "        'net_type': 'Net',\n",
        "        'no_cuda': False,\n",
        "        'seed': 1,\n",
        "        'log_interval': 100\n",
        "        }\n",
        "\n",
        "lr_anneal = [False, True]\n",
        "loss_dict = {}\n",
        "\n",
        "for i in range(len(lr_anneal)):\n",
        "    args['scheduler'] = lr_anneal[i]\n",
        "    train_loss, test_error = main(args)\n",
        "    loss_dict['train' + str(lr_anneal[i])] = train_loss\n",
        "    loss_dict['test' + str(lr_anneal[i])] = test_error\n",
        "\n",
        "with plt.xkcd():\n",
        "    fig, axs = plt.subplots(1, 2, figsize=(10,4))\n",
        "    axs[0].plot(loss_dict['trainFalse'], label='SGD', color='b')\n",
        "    axs[1].plot(loss_dict['testFalse'], label='SGD', color='b', linestyle='dashed')\n",
        "    axs[0].plot(loss_dict['trainTrue'], label='Cyclic Cosine', color='r')\n",
        "    axs[1].plot(loss_dict['testTrue'], label='Cyclic Cosine', color='r', linestyle='dashed')\n",
        "    axs[0].set_title('Train')\n",
        "    axs[1].set_title('Test')\n",
        "    axs[0].set_ylabel('Loss')\n",
        "    axs[1].set_ylabel('Error')\n",
        "    axs[0].set_xlabel('Epoch')\n",
        "    axs[1].set_xlabel('Epoch')\n",
        "    axs[0].legend()\n",
        "    axs[1].legend()\n",
        "    plt.tight_layout()\n",
        "    plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wB9L3i-Y_gNj"
      },
      "source": [
        "# Part 3: Stochastic Weight Averaging (SWA) Optimizer\n",
        "\n",
        "In this section, you'll implement a new optimizer to perform digit classification in the MNIST dataset.\n",
        "\n",
        "<p align=\"center\">\n",
        "  <img width=\"800\" height=\"200\" src=\"https://miro.medium.com/max/2000/1*_USiR_z8PKaDuIcAs9xomw.png\">\n",
        "</p>\n",
        "\n",
        "The intuition for SWA comes from the empirical observation that local minima tend to accumulate at the border of areas on loss surface where the loss value is low (points $W_1$, $W_2$ and $W_3$ are at the border of the red area of low loss in the left panel of figure above). By taking the average of several such points, it is possible to achieve a wide, generalizable solution with even lower loss ($W_{SWA}$ in the left panel of the figure above).\n",
        "\n",
        "Here is how it works. Instead of an ensemble of many models, you only need two models:\n",
        "\n",
        "$$w_{SWA}\\leftarrow \\frac{w_{SWA}\\cdot n_{models}+w}{n_{models}+1}$$\n",
        "\n",
        "*   The first model stores the running average of model weights ($w_{SWA}$ in the above formula). This will be the final model after the end of the training, and will be used for predictions.\n",
        "*   The second model ($w$ in the formula) will be traversing the weight space.\n",
        "\n",
        "At the end of each training epoch, the current weights of the second model will be used to update the weight of the running average model by taking the weighted mean between the old running average weights and the new set of weights from the second model.(See formula above.) By following this approach, we only need to train one model, and store only two models in memory during training. For prediction, we only need the running average model.\n",
        "\n",
        "Here is the pseudo code of the algorithm:\n",
        "  ```python\n",
        "    # Algorithm: Stochastic Weight Averaging\n",
        "\n",
        "    Require: Weight w_hat, learning rate lr, epochs n\n",
        "    Initialize with w <- w_hat and w_swa <- w\n",
        "    for i = 1, 2,..., n do\n",
        "        w <- w - lr * d(L_i)/dw\n",
        "        n_models <- i\n",
        "        w_swa <- (w_swa * n_models + w) / (n_models + 1)\n",
        "    end for\n",
        "    Output: w_swa\n",
        "  ```\n",
        "<p align=\"center\">\n",
        "  <img width=\"450\" height=\"250\" src=\"https://pytorch.org/assets/images/swapytorch4.png\">\n",
        "</p>\n",
        "\n",
        "We expect solutions that are centered in the flat region of the loss to generalize better than those near the boundary. Indeed, train and test error surfaces are not perfectly aligned in the weight space. Solutions that are centered in the flat region are not as susceptible to the shifts between train and test error surfaces as those near the boundary. In the figure above, the train and test error surfaces along the direction connecting the SWA and SGD solutions. As you can see, while the SWA solution has a higher training loss than the SGD solution, it is centered in a region of low loss and has a substantially better test error. To see more details about SWA, check the [paper](https://arxiv.org/pdf/1803.05407.pdf).\n",
        "  "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0Hm3y_ZK3tSH"
      },
      "source": [
        "You will implement SWA to perform digit classification. Towards this goal, you have to perform the following steps:\n",
        "\n",
        "1. Define a neural network to digit classification (provided)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BCV-Cp9GwDfG",
        "cellView": "form"
      },
      "source": [
        "# @title Network Class (run me)\n",
        "class Net(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(Net, self).__init__()\n",
        "        self.dropout = nn.Dropout(0.5)\n",
        "        self.fc1 = nn.Linear(784, 128)\n",
        "        self.fc2 = nn.Linear(128, 10)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = torch.flatten(x, 1)\n",
        "        x = self.fc1(x)\n",
        "        x = F.relu(x)\n",
        "        x = self.dropout(x)\n",
        "        x = self.fc2(x)\n",
        "        output = F.log_softmax(x, dim=1)\n",
        "        return output"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HjPY0Hjf4H-c"
      },
      "source": [
        "2. Implement training and testing procedures (provided)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "T2H0rhdpwWmS",
        "cellView": "form"
      },
      "source": [
        "# @title Train and Test functions (run me)\n",
        "def train(args, model, device, train_loader, optimizer, epoch):\n",
        "    model.train()\n",
        "    avg_loss = 0.\n",
        "    for batch_idx, (data, target) in enumerate(train_loader):\n",
        "        data, target = data.to(device), target.to(device)\n",
        "        optimizer.zero_grad()\n",
        "        output = model(data)\n",
        "        loss = F.nll_loss(output, target)\n",
        "        avg_loss += loss.item()\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "        if batch_idx % args['log_interval'] == 0:\n",
        "            print('Train Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}'.format(\n",
        "                epoch, batch_idx * len(data), len(train_loader.dataset),\n",
        "                100. * batch_idx / len(train_loader), loss.item()))\n",
        "    avg_loss /= len(train_loader.dataset)\n",
        "    return avg_loss\n",
        "\n",
        "def test(model, device, test_loader):\n",
        "    model.eval()\n",
        "    test_loss = 0\n",
        "    correct = 0\n",
        "    with torch.no_grad():\n",
        "        for data, target in test_loader:\n",
        "            data, target = data.to(device), target.to(device)\n",
        "            output = model(data)\n",
        "            test_loss += F.nll_loss(output, target, reduction='sum').item()  # sum up batch loss\n",
        "            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability\n",
        "            correct += pred.eq(target.view_as(pred)).sum().item()\n",
        "\n",
        "    test_loss /= len(test_loader.dataset)\n",
        "\n",
        "    print('\\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.4f}%)\\n'.format(\n",
        "        test_loss, correct, len(test_loader.dataset),\n",
        "        100. * correct / len(test_loader.dataset)))\n",
        "    return 1. - correct / len(test_loader.dataset)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "weKGE_hy4X2e"
      },
      "source": [
        "3. Define the SWA moving average function. Insert your code in the cell below."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "myNgg9y5xl2q"
      },
      "source": [
        "def moving_average(net1, net2, swa_n):\n",
        "    for param1, param2 in zip(net1.parameters(), net2.parameters()):\n",
        "        ####################################################################\n",
        "        # Fill in missing code below (...),\n",
        "        # then remove or comment the line below to test your function\n",
        "        raise NotImplementedError(\"Add the missing code\")\n",
        "        ####################################################################  \n",
        "        \n",
        "\n",
        "        \n",
        "        ####################################################################  "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AvyPrirp4w0z"
      },
      "source": [
        "4. Load the MNIST dataset from torchvision, initialize models, optimizers, and learning rate schedueler (provided).\n",
        "\n",
        "5. Finish the SWA update rule. Insert your code in the cell below."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tFAiwKNywc0M"
      },
      "source": [
        "def main(args):\n",
        "    # gpu configuration\n",
        "    use_cuda = not args['no_cuda'] and torch.cuda.is_available()\n",
        "    torch.manual_seed(args['seed'])\n",
        "    device = torch.device('cuda' if use_cuda else 'cpu')\n",
        "    # data loader configuration\n",
        "    train_kwargs = {'batch_size': args['batch_size']}\n",
        "    test_kwargs = {'batch_size': args['test_batch_size']}\n",
        "    if use_cuda:\n",
        "        cuda_kwargs = {'num_workers': 1,\n",
        "                       'pin_memory': True,\n",
        "                       'shuffle': True}\n",
        "        train_kwargs.update(cuda_kwargs)\n",
        "        test_kwargs.update(cuda_kwargs)\n",
        "    # load MNIST dataset\n",
        "    transform=transforms.Compose([\n",
        "        transforms.ToTensor(),\n",
        "        transforms.Normalize((0.1307,), (0.3081,))\n",
        "        ])\n",
        "    train_loader = torch.utils.data.DataLoader(datasets.MNIST('../data', train=True, download=True,\n",
        "                       transform=transform),**train_kwargs)\n",
        "    test_loader = torch.utils.data.DataLoader(datasets.MNIST('../data', train=False,\n",
        "                       transform=transform), **test_kwargs)\n",
        "    # init models\n",
        "    model = Net().to(device)\n",
        "    swa_model = Net().to(device) # new introduced model!\n",
        "    swa_n = 0\n",
        "    # config optimizer and learnng rate scheduler\n",
        "    optimizer = optim.SGD(model.parameters(), lr=args['lr'], momentum=args['momentum'])\n",
        "    scheduler = StepLR(optimizer, step_size=1, gamma=args['gamma'])\n",
        "    # training and testing\n",
        "    train_list, test_list, swa_list = [], [], []\n",
        "    for epoch in range(1, args['epochs'] + 1):\n",
        "        loss = train(args, model, device, train_loader, optimizer, epoch) # train\n",
        "        train_list.append(loss)\n",
        "        error = test(model, device, test_loader) # test\n",
        "        test_list.append(error)\n",
        "        # swa update rule\n",
        "        if epoch >= args['swa_start']:\n",
        "            ####################################################################\n",
        "            # Fill in missing code below (...),\n",
        "            # then remove or comment the line below to test your function\n",
        "            raise NotImplementedError(\"Add the missing code\")\n",
        "            ####################################################################  \n",
        "            \n",
        "\n",
        "            \n",
        "            ####################################################################\n",
        "        swa_error = test(swa_model, device, test_loader) # test swa\n",
        "        swa_list.append(swa_error)\n",
        "        scheduler.step()\n",
        "\n",
        "    return train_list, test_list, swa_list"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6ab1tpvu5mAP"
      },
      "source": [
        "6. Starting training and testing. Then plot the training loss and testing error comparison between SGD and SWA using the code provided.\n",
        "  **Note**: we might not be able to observe an expected performance of SWA, since MNIST calssification is relatively a simple task compared to CIFAR-10 / CIFAR-100. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_H7-WOhw1HWS",
        "cellView": "form"
      },
      "source": [
        "# @title Train and Plot (run me)\n",
        "# Training settings\n",
        "args = {'batch_size': 64,\n",
        "        'test_batch_size': 1000,\n",
        "        'epochs': 20,\n",
        "        'swa_start': 10,\n",
        "        'lr': 0.01,\n",
        "        'momentum': 0.9,\n",
        "        'gamma': 0.7,\n",
        "        'net_type': 'Net',\n",
        "        'no_cuda': False,\n",
        "        'seed': 1,\n",
        "        'log_interval': 100\n",
        "        }\n",
        "\n",
        "train_loss, test_error, swa_error = main(args)\n",
        "\n",
        "with plt.xkcd():\n",
        "    plt.plot(train_loss, label='Train Loss', linewidth=2)\n",
        "    plt.plot(test_error, label='SGD Test Error', linewidth=2)\n",
        "    plt.plot(range(10, 20), swa_error[10:], label='SWA Test Error', linewidth=2)\n",
        "    plt.xlabel('Epoch')\n",
        "    plt.legend()\n",
        "    plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VLA3CSrQauwt"
      },
      "source": [
        "# Part 4: Better-Know-A-Pod\n",
        "*This section is designed to let you better know your pod mates. You are encouraged to discuss these with your pod mates during pod sessions or during pod social events, but NOT over Slack. The idea is to get to know your pod mates face-to-face.*\n",
        "\n",
        "**List 2 other members of your pod and write a thoughtful and nice short blurb about them (around 100 words each). Where are they from? What do they study? What are their hobbies? What do you have in common?**\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "utlLXvQDb9R5"
      },
      "source": [
        "better_know_a_pod = '''\n",
        "Your answer to Better-Know-A-Pod Here (100 words each).\n",
        "'''"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7J3FEikx_xR0"
      },
      "source": [
        "# Submission\n",
        "Double check that the answers have been correctly formatted in the Airtable below, and edit them as desired. **In addition to submitting the form, contribute and reflect on these readings in the student Slack space.**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EZvmWNqW_YTB",
        "cellView": "form"
      },
      "source": [
        "#@markdown #Run Cell to Show Airtable Form\n",
        "#@markdown ##**Confirm your answers and then click \"Submit\"**\n",
        "import numpy as np\n",
        "import time\n",
        "from IPython.display import IFrame\n",
        "def prefill_form(src, fields: dict):\n",
        "  '''\n",
        "  src: the original src url to embed the form\n",
        "  fields: a dictionary of field:value pairs,\n",
        "  e.g. {\"pennkey\": my_pennkey, \"location\": my_location}\n",
        "  '''\n",
        "  prefills = \"&\".join([\"prefill_%s=%s\"%(key, fields[key]) for key in fields])\n",
        "  src = src + prefills\n",
        "  src = \"+\".join(src.split(\" \"))\n",
        "  return src\n",
        "\n",
        "#autofill fields if they are not present\n",
        "#a missing pennkey and pod will result in an Airtable warning\n",
        "#which is easily fixed user-side.\n",
        "try: my_pennkey;\n",
        "except NameError: my_pennkey = \"\"\n",
        "try: my_pod;\n",
        "except NameError: my_pod = \"Select\"\n",
        "try: rephrase_1;\n",
        "except NameError: rephrase_1 = \"\"\n",
        "try: rephrase_2;\n",
        "except NameError: rephrase_2 = \"\"\n",
        "try: rephrase_3;\n",
        "except NameError: rephrase_3 = \"\"\n",
        "try: rephrase_4;\n",
        "except NameError: rephrase_4 = \"\"\n",
        "try: bound;\n",
        "except NameError: bound = \"\"\n",
        "try: bias_correction;\n",
        "except NameError: bias_correction = \"\"\n",
        "try: regret;\n",
        "except NameError: regret = \"\"\n",
        "try: better_know_a_pod;\n",
        "except NameError: better_know_a_pod = \"\"\n",
        "\n",
        "\n",
        "#autofill time if it is not present\n",
        "try: t0;\n",
        "except NameError: t0 = time.time()\n",
        "try: t1;\n",
        "except NameError: t1 = time.time()\n",
        "times = np.array([t1])-t0\n",
        "\n",
        "fields = {\"pennkey\": my_pennkey,\n",
        "          \"pod\": my_pod,\n",
        "          \"rephrase_1\":rephrase_1,\n",
        "          \"rephrase_2\":rephrase_2,\n",
        "          \"rephrase_3\":rephrase_3,\n",
        "          \"rephrase_4\":rephrase_4,\n",
        "          \"bound\":bound,\n",
        "          \"bias_correction\":bias_correction,\n",
        "          \"regret\":regret,\n",
        "          \"better_know_a_pod\": better_know_a_pod,\n",
        "          \"cumulative_times\": times}\n",
        "\n",
        "src = \"https://airtable.com/embed/shrN8MK580CBFiGCu?\"\n",
        "\n",
        "#now instead of the original source url, we do: src = prefill_form(src, fields)\n",
        "display(IFrame(src = prefill_form(src, fields), width = 800, height = 400))"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}